import math
import random

class KMeans:
    def __init__(self, n_clusters=2, max_iter=100, random_state=42):
        self.n_clusters = n_clusters
        self.max_iter = max_iter
        self.random_state = random_state
        self.centroids = []
        self.labels = []

    def _euclidean_distance(self, x1, x2):
        return math.sqrt(sum((a - b)**2 for a, b in zip(x1, x2)))

    def _initialize_centroids(self, X):
        random.seed(self.random_state)
        self.centroids = random.sample(X, self.n_clusters)

    def _assign_clusters(self, X):
        labels = []
        for x in X:
            distances = [self._euclidean_distance(x, centroid) for centroid in self.centroids]
            labels.append(distances.index(min(distances)))
        return labels

    def _update_centroids(self, X, labels):
        clusters = [[] for _ in range(self.n_clusters)]
        for x, label in zip(X, labels):
            clusters[label].append(x)
        
        new_centroids = []
        for cluster in clusters:
            if not cluster:
                new_centroids.append(random.choice(X))
                continue
            dim = len(cluster[0])
            centroid = [sum(x[i] for x in cluster) / len(cluster) for i in range(dim)]
            new_centroids.append(centroid)
        return new_centroids

    def fit(self, X):
        self._initialize_centroids(X)
        for _ in range(self.max_iter):
            self.labels = self._assign_clusters(X)
            old_centroids = self.centroids.copy()
            self.centroids = self._update_centroids(X, self.labels)
            if all(self._euclidean_distance(old, new) < 1e-6 for old, new in zip(old_centroids, self.centroids)):
                break
        print("K-Means聚类完成")

    def predict(self, X):
        return [self._assign_clusters([x])[0] for x in X]

if __name__ == "__main__":
    iris_data = [
        [5.1, 3.5], [4.9, 3.0], [4.7, 3.2], [4.6, 3.1], [5.0, 3.6],
        [5.4, 3.9], [4.6, 3.4], [5.0, 3.4], [4.4, 2.9], [4.9, 3.1],
        [5.4, 3.7], [4.8, 3.4], [4.8, 3.0], [4.3, 3.0], [5.8, 4.0],
        [5.7, 4.4], [5.4, 3.9], [5.1, 3.5], [5.7, 3.8], [5.1, 3.8],
        [7.0, 3.2], [6.4, 3.2], [6.9, 3.1], [5.5, 2.3], [6.5, 2.8],
        [5.7, 2.8], [6.3, 3.3], [4.9, 2.4], [6.6, 2.9], [5.2, 2.7],
        [6.3, 2.5], [6.0, 2.2], [6.1, 2.8], [5.6, 2.9], [5.5, 2.5],
        [6.1, 3.0], [5.8, 2.6], [5.0, 2.3], [5.6, 2.7], [5.7, 3.0],
        [6.7, 3.1], [6.3, 2.3], [5.8, 2.7], [6.5, 3.0], [5.1, 2.5],
        [5.8, 2.8], [5.9, 3.0], [5.0, 3.0], [5.5, 3.5], [5.4, 3.0]
    ]

    kmeans = KMeans(n_clusters=3, max_iter=50, random_state=42)
    kmeans.fit(iris_data)

    print(f"聚类中心：{kmeans.centroids}")
    print(f"前10个样本的簇标签：{kmeans.labels[:10]}")
